# \author  Neil Turton <neilt@amd.com>
#  \brief  Nanotube testing
#   \date  2020-08-03
#
###########################################################################
# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
# SPDX-License-Identifier: MIT
###########################################################################
Import('env')

import glob
import os.path
import subprocess

# Read the test list.  The generating script produces a list of tests
# on stdout, one per line.  Each test name consists of a number of
# fields separated by commas.  The first field is the base name.  The
# remaining fields are pass names.
def run_gen_test_list(test_list, dir, args):
    proc = subprocess.Popen([os.path.join(str(dir), "gen_test_list")] + args,
                            stdout=subprocess.PIPE)
    for line in proc.stdout:
        test_list.append(line.decode('utf-8').rstrip().split(","))
    proc.wait()
    assert proc.returncode == 0, proc.returncode

# Translate pass names to the right command line parameter
pass_opts = {
    'converge': (
        '-move-alloca -compact-geps '
        '-converge_mapa'
    ),
    'byteify' : '-byteify',
    'destruct': '-destruct',
    'ebpf2nt': '-codegenprepare -instsimplify -dce -ebpf2nanotube',
    'mem2req' : '-instsimplify -mem2req',
    'platform' : '-platform -always-inline -instsimplify',
    'control' : '-control-capsule -always-inline',
    'ntattr' : ( '-nt-attributes -O2' ),
    'optreq' :  (
        '-mergereturn '
        '-optreq -enable-loop-unroll -always-inline '
        '-instsimplify -loop-unroll -simplifycfg' ),
    'inline'  : ( '-always-inline -instsimplify' ),
    'inline_opt': ( '-always-inline'
                    ' -rewrite-setup -replace-malloc'
                    ' -thread-const -instsimplify'
                    ' -codegenprepare -instsimplify -dce '
                    ' -enable-loop-unroll -loop-unroll'
                    ' -move-alloca -simplifycfg -instcombine'
                    ' -thread-const -instsimplify -simplifycfg' ),
    'pipeline' : ( '-compact-geps -basic-aa -tbaa '
                   '-nanotube-aa -pipeline' ),
    'flatten' : ( '-flatten-cfg' ),
}

# Additional build steps for the tests
def link_taps(env, out_bc_file, in_bc_file):
    env.LlvmLink([out_bc_file],
                 [
                     in_bc_file,
                     env['BUILD_TOP'].File('libnt/nanotube_low_level.bc'),
                 ])

def lower_api(env, out_bc_file, in_bc_file):
    env.LlvmLink([out_bc_file],
                 [
                     in_bc_file,
                     env['BUILD_TOP'].File('libnt/nanotube_high_level.bc'),
                 ])

def build_hls_exe(env, exe_file, in_bc_file):
    json_file = env.Nanotube_Be(in_bc_file)
    out = env.Build_Kernel_Test(exe_file, json_file)

# Adds the test with build instructions to the target
# test_info: ("src:opts", pass1, pass2, pass3, ...)
def add_test(test_info):
    src_file_and_opts = test_info[0].split(':')
    src_file = src_file_and_opts[0]
    extra_opts = src_file_and_opts[1]
    passes = list(test_info[1:])
    base  = os.path.splitext(src_file)[0]

    out_pass_str = ("".join("."+x for x in passes))
    out_bc_file = base + out_pass_str + ".bc"
    exe_file = "test_" + base + out_pass_str

    if len(passes) == 0:
        env.BitFile(src_file)
    else:
        this_pass = passes[-1]
        in_pass_str = ("".join("."+x for x in passes[:-1]))
        in_bc_file = base + in_pass_str + ".bc"

        if this_pass == "hls":
            build_hls_exe(env, exe_file, in_bc_file)
            return

        if this_pass == "link_taps":
            link_taps(env, out_bc_file, in_bc_file)
        elif this_pass == "lower":
            lower_api(env, out_bc_file, in_bc_file)
        else:
            env.NanotubeOpt(out_bc_file, in_bc_file,
                            OPT_FLAGS=pass_opts[this_pass]+" "+extra_opts)

    obj_file = env.BitFileObj(out_bc_file)
    env.Program(
        exe_file,
        [
            obj_file,
        ]
    )

# Go through the test list and create rules to build intermediate build
# artifacts, making sure to remove duplicates
#
# Example:
#   foo,pass_A, pass_B
#   foo,pass_A, pass_X, pass_Y
# will add foo.pass_A.pass_B, foo.pass_A, foo, foo.pass_A.pass_X.pass_Y,
# foo.pass_A.pass_X in that sequence
def test_rules(test_list):
    done = set()
    for test_info in test_list:
        test_info = tuple(test_info)
        while test_info not in done:
            done.add(test_info)
            add_test(test_info)
            if len(test_info) <= 1:
                break
            test_info = test_info[:-1]

# Generate rules for golden input / output data
# Input / output data is same for various compilation variants of the same test
# so remove duplicates after stripping different build options from all the
# tests
def golden_rules(test_list, golden_src):
    data_prefixes_with_dups = [ti[0].split('.')[0] for ti in test_list]
    data_prefixes = sorted(set(data_prefixes_with_dups))
    for base in data_prefixes:
        for direction in ('IN', 'OUT'):
            out_file = "golden/test_"+base+".pcap."+direction
            in_file = golden_src.File("test_"+base+".packets."+direction)
            env.Command(out_file,
                        [env['SOURCE_TOP'].File('scripts/text_to_pcap'), in_file],
                        [ "${SOURCES[0]} ${SOURCES[1]} ${TARGETS[0]}" ])

env.add_libs('test_harness')

env.Append(CCFLAGS='-Wno-gcc-compat'\
                   ''.split())

# Build all tests and golden input / output for a specific directory
def build_tests_and_golden(source_dir):
    golden_src = source_dir.Dir('golden')
    test_list = []
    run_gen_test_list(test_list, source_dir, [])
    test_rules(test_list)
    golden_rules(test_list, golden_src)

build_tests_and_golden(env['SOURCE_DIR'])

# Allow the user to specify additional tests
extra = env.get('EXTRA_TESTS')
if extra:
    extra_source = Dir(extra).Dir("testing/kernel_tests")
    build_tests_and_golden(extra_source)
    # Build the extra tests into the same build directory
    env.VariantDir(env['BUILD_DIR'], extra_source, duplicate=0)
    # Additional includes from the extra test directory
    env.Append(CCFLAGS=["-I", extra_source.Dir("include")])
    env.Append(CCFLAGS=["-I", Dir(extra).Dir("include")])
    # Add any extra libraries
    extra_lib = env.get('EXTRA_LIBS')
    if extra_lib:
        env.Append(LIBPATH=[Dir(extra).Dir("build")])
        env.Append(LIBS=[extra_lib])
        env.Append(RPATH=[Dir(extra).Dir("build")])


llvm_components = env.get('EXTRA_LLVM_COMPONENTS')
if llvm_components:
    llvm_config = env.get('LLVM_CONFIG')
    args = [llvm_config, '--libs']
    args.extend(llvm_components.split())
    llvm_libs = subprocess.check_output(args, universal_newlines=True).rstrip('\n')

    env.Append(RPATH=[env['LLVM_LIB_DIR']])
    env.Append(LIBPATH=[env['LLVM_LIB_DIR']])
    env.Append(LIBS=llvm_libs.split())
